import random
import networkx as nx

# 创建一个网络图
G = nx.Graph()


def update_traffic_matrix(num_tor):
    """生成一个示例的流量需求矩阵（模拟变化）"""
    traffic_matrix = {}
    for i in range(1, num_tor + 1):
        traffic_matrix[i] = {}
        for j in range(1, num_tor + 1):
            if i != j:
                traffic_matrix[i][j] = random.randint(1, 100)  # 模拟随机的流量需求，单位Mbps
            else:
                traffic_matrix[i][j] = 0  # 自己与自己的流量需求为0
    return traffic_matrix


def adjust_optical_links(G, traffic_matrix, threshold=50):
    """
    根据流量需求矩阵调整光交换机链路带宽。
    如果流量需求超过阈值，则增加光链路带宽。
    """
    optical_switches = [n for n in G.nodes if G.nodes[n].get('type') == 'optical_switch']
    tor_switches = [n for n in G.nodes if G.nodes[n].get('type') == 'tor']

    for tor_id in tor_switches:
        for sw_node in optical_switches:
            # 检查当前ToR和光交换机之间的流量需求
            if tor_id != sw_node:
                traffic_demand = traffic_matrix.get(tor_id, {}).get(sw_node, 0)
                if traffic_demand > threshold:
                    # 如果流量需求超过阈值，增加光链路带宽
                    if G.has_edge(tor_id, sw_node):
                        G[tor_id][sw_node]['capacity'] = max(G[tor_id][sw_node]['capacity'], traffic_demand)
                    else:
                        # 如果还没有建立链路，创建新的链路并设置带宽
                        G.add_edge(tor_id, sw_node, type='optical', capacity=traffic_demand, port_type='optical')
                    print(f"增加了{tor_id}与{sw_node}之间的光链路带宽，带宽：{traffic_demand} Mbps")
                else:
                    # 如果流量需求低于阈值，则不做调整
                    if G.has_edge(tor_id, sw_node):
                        G[tor_id][sw_node]['capacity'] = min(G[tor_id][sw_node]['capacity'], threshold)
                    print(f"流量需求低，{tor_id}与{sw_node}之间保持默认带宽")


# 生成流量需求矩阵
num_tor = 16
traffic_matrix = update_traffic_matrix(num_tor)

# 向图中添加光交换机和ToR交换机节点
for i in range(1, num_tor + 1):
    G.add_node(i, type='tor')

# 添加光交换机节点
G.add_node(100, type='optical_switch')

# 调用方法，调整光链路
adjust_optical_links(G, traffic_matrix)
